from absl import app, flags
import datetime
import pytz
import pandas as pd
import numpy as np
import functools
import traceback
from concurrent.futures import ProcessPoolExecutor

from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.tool.strat_monitor.util.monitor_util import get_topic_from_db
from coin.util.queue.tools.kafka_archive import run_from_kafka_archive
from coin.util.queue.config import KafkaConfig
from coin.support.proto_log.logic.util import run_from_strat_log_archive

FLAGS = flags.FLAGS

format_ts = lambda x: "%.2f" % x if isinstance(x, float) else (x if isinstance(x, str) else '')


def get_ts_from_str(dt):
  dt = datetime.datetime.strptime(dt, "%Y-%m-%d %H:%M:%S")
  dt = dt.replace(tzinfo=pytz.UTC)
  ts = dt.timestamp() * 1e+9
  return ts


class StrategyLogHandle(object):
  def __init__(self):
    self.config_log = []
    self.pnl_balance = []
    self.order_event = {}
    return

  def handle_for_pnlbalace(self, ts, pb):
    if not self.pnl_balance:
      self.pnl_balance.append(ts)
    elif abs(ts - self.pnl_balance[-1]) < 1e+9:
      self.pnl_balance[-1] = ts
    else:
      self.pnl_balance.append(ts)

  def handle_for_og(self, ts, pb):
    og_log = pb.og_log
    if og_log.type == OrderGatewayLog.ORDER_EVENT:
      key = '_'.join([str(og_log.event.order_event_id), str(og_log.event.event_time)])
      self.order_event[key] = ts

  def handle_for_config(self, ts, config_log):
    if not self.config_log:
      self.config_log.append(ts)
    elif abs(ts - self.config_log[-1]) < 1e+9:
      self.config_log[-1] = ts
    else:
      self.config_log.append(ts)

  def handle(self, partition, record):
    topic = record.topic
    ts = int(record.timestamp * 1e+6)
    pb = StrategyLog()
    pb.ParseFromString(record.value)
    if pb.HasField('pnl_balance'):
      self.handle_for_pnlbalace(ts, pb)
    if pb.type == StrategyLog.OG_LOG:
      self.handle_for_og(ts, pb)
    if pb.type == StrategyLog.PASS_EXECUTOR_CONFIG:
      self.handle_for_config(ts, pb.pass_executor_config)
    if pb.type == StrategyLog.AGG_EXEUCTOR_CONFIG:
      self.handle_for_config(ts, pb.agg_executor_config)

  def on_log(self, timestamp, log):
    pb = StrategyLog()
    pb.ParseFromString(log)
    if pb.HasField('pnl_balance'):
      self.handle_for_pnlbalace(timestamp, pb)
    if pb.type == StrategyLog.OG_LOG:
      self.handle_for_og(timestamp, pb)


def compare_ts_diff(log_from_kafka: dict, log_from_local: dict):
  res = {}
  for key in log_from_kafka:
    ts_from_kafka = log_from_kafka[key]
    ts_from_local = log_from_local.get(key)
    if ts_from_local:
      res[ts_from_kafka] = (ts_from_kafka - ts_from_local) / 1e+9
  return res


def guess_ts_diff(log_from_kafka: list, log_from_local: list):
  res = {}
  if len(log_from_kafka) < 10 or len(log_from_local) < 10:
    return res
  for ts_from_kafka in log_from_kafka:
    diff = min([abs(ts_from_kafka - ts) for ts in log_from_local])
    res[ts_from_kafka] = diff / 1e+9
  return res


def diff_result_analytic(dif_res: dict):
  dif = pd.Series(dif_res)
  describe = dif.describe(percentiles=[0.5, 0.99])
  time_window = 300 * 1e+9
  extre_dif = dif[dif > 30]
  extre_ts_list = []
  for ts in extre_dif.index:
    extre_ts_list.append([ts, dif.loc[ts:ts + time_window].mean()])
  extre_ts_list.sort(key=lambda x: x[-1])
  if extre_ts_list:
    describe.loc['extre_5m_avg'] = extre_ts_list[-1][1]
  else:
    describe.loc['extre_5m_avg'] = None
  describe = describe.drop(index=['std', 'min'])
  describe = describe.replace(np.nan, '')
  describe = describe.apply(format_ts)
  return describe


def time_diff_analytic(start_ts, end_ts, topic):
  handler_from_kafka = StrategyLogHandle()
  kafka_config = KafkaConfig.from_cmd_config(FLAGS.kafka_config_filename)
  run_from_kafka_archive(topics=[topic],
                         kafka_config=kafka_config,
                         timestamp_from=int(start_ts),
                         timestamp_to=int(end_ts),
                         callback=handler_from_kafka.handle)
  strat_log_root_dir = FLAGS.root_dir or '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log/'
  handler_from_local = StrategyLogHandle()
  _, machine, strategy_name = topic.split('_', 2)
  run_from_strat_log_archive(on_log_callback=handler_from_local.on_log,
                             start_time=datetime.datetime.utcfromtimestamp(start_ts / 1e+9),
                             end_time=datetime.datetime.utcfromtimestamp(end_ts / 1e+9),
                             root_dir=strat_log_root_dir,
                             machine=machine,
                             strategy_name=strategy_name)
  order_event_dif = compare_ts_diff(handler_from_kafka.order_event, handler_from_local.order_event)
  pnl_balance_dif = guess_ts_diff(handler_from_kafka.pnl_balance, handler_from_local.pnl_balance)
  order_event_analytic = diff_result_analytic(order_event_dif)
  pnl_balance_analytic = diff_result_analytic(pnl_balance_dif)
  cols = [('order_event', idx) for idx in order_event_analytic.index
         ] + [('pnl_balance', idx) for idx in pnl_balance_analytic.index]
  res = pd.DataFrame([order_event_analytic.tolist() + pnl_balance_analytic.tolist()])
  res.columns = pd.MultiIndex.from_tuples(cols)
  res.index = pd.MultiIndex.from_tuples([(machine, strategy_name)])
  return res


def main(_):
  pd.set_option('display.width', 300)
  pd.set_option('max_colwidth', 300)
  start_ts = get_ts_from_str(FLAGS.start_dt)
  end_ts = get_ts_from_str(FLAGS.end_dt)
  topics = FLAGS.topics or get_topic_from_db(FLAGS.kafka_config_filename, active=True)
  if isinstance(topics, str):
    topics = topics.split(',')
  result = pd.DataFrame()
  func_list = []
  for topic in topics:
    analytic_func = functools.partial(time_diff_analytic,
                                      start_ts=start_ts,
                                      end_ts=end_ts,
                                      topic=topic)
    func_list.append((topic, analytic_func))
  with ProcessPoolExecutor(max_workers=len(func_list)) as executor:
    future_list = [(topic, executor.submit(func)) for topic, func in func_list]
  for topic, future in future_list:
    try:
      result = result.append(future.result())
    except Exception as e:
      print('Fail to analytic strategy log.', topic, type(e), e)
      print(traceback.format_exc())

  print(result.to_string())


if __name__ == '__main__':
  flags.DEFINE_string('start_dt', None, '%Y-%m-%d %H:%M:%S')
  flags.DEFINE_string('end_dt', None, '%Y-%m-%d %H:%M:%S')
  flags.DEFINE_string('topics', None, 'topics')
  flags.DEFINE_string('kafka_config_filename',
                      '../../coin_deploy/support_monitor/config/kafka_config.json', 'kafka config')
  flags.DEFINE_string('root_dir', None, 'root_dir')
  app.run(main)
